use std::time::Instant;

use anyhow::{anyhow, Result};
use futures_util::StreamExt;
use log::debug;
use whisper_rs::{FullParams, SamplingStrategy, WhisperContext, WhisperState, WhisperToken};

use crate::domain::{Segment, PCM_SIZE_ONE_SECOND};

#[derive(Copy, Clone)]
pub struct TranscriberConfig {
    step_ms: u32,
    min_samples_size: usize,
    n_segments_cache: usize,
    max_prompt_token: usize,
}

impl Default for TranscriberConfig {
    fn default() -> Self {
        Self {
            step_ms: 3000,
            n_segments_cache: 2,
            min_samples_size: PCM_SIZE_ONE_SECOND * 5,
            max_prompt_token: 128,
        }
    }
}

pub struct Transcriber<'a> {
    pcm: Vec<f32>,
    cached_offset: usize,
    state: WhisperState<'a>,
    config: TranscriberConfig,
    preset_tokens: Vec<WhisperToken>,
    prompt_tokens: Vec<WhisperToken>,
}

impl<'a> Transcriber<'a> {
    pub fn new(ctx: &'a WhisperContext) -> Result<Self> {
        let config = TranscriberConfig::default();

        Ok(Self {
            config,
            cached_offset: 0,
            preset_tokens: Vec::new(),
            state: ctx.create_state()?,
            pcm: Vec::with_capacity(30 * PCM_SIZE_ONE_SECOND as usize),
            prompt_tokens: Vec::with_capacity(config.max_prompt_token),
        })
    }

    pub fn try_inference(&mut self, pcm: &[f32]) -> Result<Option<Vec<Segment>>> {
        self.pcm.extend_from_slice(pcm);
        if self.pcm.len() - self.cached_offset < self.config.min_samples_size {
            return Ok(None);
        }

        println!("pcm len: {}", self.pcm.len());

        let mut params = FullParams::new(SamplingStrategy::Greedy { best_of: 0 });
        params.set_language(Some("en"));
        params.set_n_threads(num_cpus::get() as _);
        let tokens = [&self.preset_tokens[..], &self.prompt_tokens[..]].concat();
        if !tokens.is_empty() {
            params.set_tokens(&tokens);
        }

        let start = Instant::now();
        self.state.full(params, &self.pcm[..])?;

        let n_segments = self.state.full_n_segments()?;
        if n_segments == 0 {
            return Ok(None);
        }

        let elapsed_mills = start.elapsed().as_millis();
        if elapsed_mills > self.config.step_ms.into() {
            debug!(
                "whisper full took {}ms for {} size pcm.",
                elapsed_mills,
                self.pcm.len()
            );
            // return Err(anyhow!(
            //     "whisper full took {} ms too slow than {}ms for {} size pcm",
            //     elapsed_mills,
            //     self.config.step_ms,
            //     self.pcm_f32.len()
            // ));
        }

        let mut prompt_tokens = Vec::with_capacity(n_segments as usize);
        let mut segments = Vec::with_capacity(n_segments as usize);

        for n_segment in 0..n_segments {
            let start = self.state.full_get_segment_t0(n_segment)?;
            let end = self.state.full_get_segment_t1(n_segment)?;
            let text = self.state.full_get_segment_text(n_segment)?;
            let segment = Segment::from((text, start, end));
            segments.push(segment);

            let num_tokens = self.state.full_n_tokens(n_segment)?;
            let mut tokens = Vec::with_capacity(num_tokens as usize);
            for n_token in 0..num_tokens {
                tokens.push(self.state.full_get_token_id(n_segment, n_token)?);
            }
            prompt_tokens.push(tokens);
        }
        assert!(!segments.is_empty());

        segments = segments
            .into_iter()
            .skip_while(|segment| {
                println!(
                    "cached_offset: {}, segment offset size: {}",
                    self.cached_offset,
                    segment.offset_size()
                );
                self.cached_offset != 0 && self.cached_offset <= segment.offset_size()
            })
            .collect();

        if segments.is_empty() {
            return Ok(None);
        }

        let new_offset = segments
            .iter()
            .take(self.config.n_segments_cache)
            .last()
            .map(|segment| segment.offset_size())
            .expect("At least one segment.");

        println!(
            "new_offset: {}, cached_offset: {}",
            new_offset, self.cached_offset
        );

        if self.cached_offset != 0 {
            let drop_size = new_offset - self.cached_offset;
            self.pcm.drain(0..drop_size);
        }
        self.cached_offset = new_offset;
        self.prompt_tokens
            .extend(prompt_tokens.into_iter().flatten());

        self.preset_tokens.drain(
            0..i32::max(
                0,
                self.preset_tokens.len() as i32 - self.config.max_prompt_token as i32,
            ) as usize,
        );

        Ok(Some(segments))
    }
}

#[cfg(test)]
mod tests {
    use hound::{WavReader, WavSpec};

    use crate::{settings::Settings, transcriber::Transcriber};

    #[test]
    fn transcriber_should_work_success() {
        let settings = Settings::new().unwrap();
        let whisper_context = settings.whisper.context().unwrap();
        let mut transcriber = Transcriber::new(&whisper_context).unwrap();

        let mut wav = WavReader::open("audio.wav").unwrap();
        let WavSpec {
            channels,
            sample_rate,
            bits_per_sample,
            ..
        } = wav.spec();
        assert_eq!(channels, 1);
        assert_eq!(sample_rate, 16000);
        assert_eq!(bits_per_sample, 16);

        let audio = whisper_rs::convert_integer_to_float_audio(
            &wav.samples::<i16>()
                .flat_map(|s| Ok::<i16, anyhow::Error>(s?))
                .collect::<Vec<_>>(),
        );

        let res = transcriber.try_inference(&audio[..]);
        println!("{res:?}");
    }
}
